
import cv2
import numpy as np
import scipy.interpolate as spi
from scipy.interpolate import UnivariateSpline
import os
import math

class keyframe():
    def __init__(self, index=None,imgrgb=None,box=None,kp=None,des=None,linepoint=None,pointcenter=None,kpall=None,desall=None):
        self.index=index#第几帧
        self.imgrgb=imgrgb
        self.box=box
        self.kp=kp#框内特征点
        self.des=des

        self.linepoint=linepoint#追踪曲线
        self.pointcenter=pointcenter#追踪曲线中点
        self.kpall = kpall#全图的特征点
        self.desall=desall

    def line_box(self,img):
        cv2.rectangle(img, (int(self.box[0, 0]), int(self.box[0, 1])), (int(self.box[3, 0]), int(self.box[3, 1])), (0, 255, 0), 2)
        return img

def keypool_update_factory():
    first_reduction_done = [False]  # 使用列表以允许在闭包内修改
    def keypool_update(kframe, k_frame, simi_frame_better,rot):
        if len(kframe) == 5 and not first_reduction_done[0]:
            del kframe[1:3]
            first_reduction_done[0] = True  # 标记第一次缩减已完成
        elif len(kframe)==16:
            kframe=kframe
        else:
            pcenter1 = simi_frame_better[0].pointcenter
            pcenter2 = k_frame.pointcenter
            pcenterdis = np.linalg.norm(np.array(pcenter1) - np.array(pcenter2))#追踪点中心点的间距
            bcenter1 = simi_frame_better[0].box[4]
            bcenter2 = k_frame.box[4]
            bcenterdis = np.linalg.norm(np.array(bcenter1) - np.array(bcenter2))#追踪框中心点的间距
            box_l = k_frame.box[5, 0]
            # threshold = box_l * 0.2
            threshold = box_l * 0.05
            if (bcenterdis > threshold and pcenterdis > threshold) or rot>8:
                kframe.append(k_frame)
        return kframe
    return keypool_update
# # 使用工厂函数创建一个新的 keypool_update 函数，它会记住 first_reduction_done 的状态
# keypool_update = keypool_update_factory()


def line_trackpoint(img,trackpoint):
    for j in range(len(trackpoint)):
        cv2.circle(img, (int(trackpoint[j][0]), int(trackpoint[j][1])), 2, (0, 225, 0), -1)
    return img

def line_track(u_f,v_f,data,img2):#用于普通帧可视化
    data_new = []
    data_newin=[]
    for j in range(len(data)):
        idx0=int(data[j][0])
        idx1=int(data[j][1])
        x=idx0 + u_f[idx1][idx0]
        y=idx1+v_f[idx1][idx0]
        data_new.append((x, y))
        if x>= 0 and x <= img2.shape[1] and y >= 0 and y<= img2.shape[0]:
            data_newin.append((x, y))
            cv2.circle(img2, (int(x), int(y)), 1, (0,255, 0), -1)
    data_new=np.array(data_new)
    return data_new,img2,data_newin


def boxuv_trackpoint(u_f,v_f,kf,box2):#box追踪的uv,用于关键帧 匹配, 由patch追踪得到uv恢复追踪点坐标
    u_f = u_f.astype(int)
    v_f = v_f.astype(int)
    center_des = box2[1, :] - kf.box[1, :]
    data_new = []
    for j in range(len(kf.linepoint)):
        idx0=int(kf.linepoint[j][0])
        idx1=int(kf.linepoint[j][1])

        x=idx0+u_f[idx1-kf.box[0,1],idx0-kf.box[0,0]]+center_des[0]
        y=idx1+v_f[idx1-kf.box[0,1],idx0-kf.box[0,0]]+center_des[1]
        data_new.append((x, y))
        # cv2.circle(img2,(x,y), 1, (0, 255, 0), -1)
    data_new = np.array(data_new)
    return data_new

def save_point(track_points,i):
    track_points = np.array(track_points)
    track_points=track_points.astype(int)
    folder="/home/ltt/point_track/trackimage1/trackingpoint/"
    file_name = f"experiment_{i}.txt"
    result_filepath = os.path.join(folder, file_name)
    np.savetxt(result_filepath, track_points)
    return track_points


def b_line(tracking):#插值
    if len(tracking)<15:
        draw_point=tracking
    else:
        tracking=tracking.astype(int)
        unique_rows = np.unique(tracking[:, 1], return_index=True, axis=0)[1]

        tracking = tracking[unique_rows]
        tracking=tracking[tracking[:, 1].argsort()]
        x_0= tracking[:,1].tolist()
        y_0= tracking[:,0].tolist()
        spline = UnivariateSpline(x_0, y_0)
        # print("spline是", spline)
        x_0min=int(x_0[0])
        x_0max=int(x_0[-1])
        x_0line = range(x_0min, x_0max+1,1)
        y_line = spline(x_0line)
        # print("y_line是",y_line)
        # print("x_line是",x_0line)

        has_nan = np.isnan(y_line)
        # 判断是否存在NaN值
        if np.any(has_nan):
        # #在y方向上插线
            track_points = tracking.copy()
            sorted_t = np.argsort(track_points[:, 1])
            # 使用索引对整个矩阵的行进行重新排序
            track_points = track_points[sorted_t]
            x = track_points[:, 1]
            y = track_points[:, 0]
            n = track_points.shape[0]
            new_points = []
            for i in range(n-1):
                l = x[i + 1] - x[i]
                if l > 0:
                    x_ = int((x[i] + x[i + 1]) / 2)
                    y_ = int((y[i] + y[i + 1]) / 2)
                    new_points.append((y_, x_))
            new_points = np.array(new_points).reshape(-1, 2)
            point = np.concatenate((track_points, new_points), axis=0)
            sorted_p = np.argsort(point[:, 0])#按x排序
            # 使用索引对整个矩阵的行进行重新排序
            points = point[sorted_p]

            #在x方向上插值
            x__=points[:,0]
            y__=points[:,1]
            k=len(points)
            x_new=[]
            for p in range(k-1):
                h=x__[p+1]-x__[p]
                if h>0:
                    x_1 = int((x__[i] + x__[i + 1]) / 2)
                    y_1 = int((y__[i] + y__[i + 1]) / 2)
                    x_new.append((x_1, y_1))
            points_f = np.array(x_new).reshape(-1, 2)
            points_f = np.concatenate((points, points_f), axis=0)
            sorted_y = np.argsort(point[:, 1])  # 按y排序
            draw_point=points_f[sorted_y].astype(np.int)

        else:
            draw_point = []
            for i in range(x_0min, x_0max + 1):
                # print("i-x_0min是", i - x_0min)
                idx = int(i - x_0min)
                y = int(y_line[idx])
                draw_point.append((y, i))
            draw_point = np.array(draw_point)
    return draw_point#x_new, y_new

